from multiprocessing import Event, Process, Condition
from multiprocessing.connection import Client
from threading import Thread


class AgentProcess(Process):
    def __init__(self, address, authkey, func):
        super().__init__()
        self.daemon = True
        self._address = address
        self._authkey = authkey
        self._conn = None
        self._cond = Condition()
        self.actions = ""
        self.observation = ""
        self.transfer_thread = None
        # self.trigger = Event()
        self._func = func
        self._game_over = False
    
    def run(self):
        self.transfer_thread = Thread(target=self.transfer_data, daemon=True)
        self.transfer_thread.start()
        while True:
            with self._cond:
                self._cond.wait()
            if self._game_over:
                break
            self.actions += self._func(1600000)  # 模拟agent引擎，生成动作
            print(f"\t  Agent -> Scheduler: {self.actions}")
            self._conn.send(self.actions)
            self.actions = ""
    
    def transfer_data(self):
        self._conn = Client(address=self._address, authkey=self._authkey)
        while True:
            rcv_data = self._conn.recv()
            print(f"\t  Agent <- Scheduler: {rcv_data}")
            if rcv_data == "ready":
                print(f"\t  Agent -> Scheduler: 'ok")
                self._conn.send("ok")
            elif rcv_data == "done":
                self._game_over = True
                with self._cond:
                    self._cond.notify_all()
                break
            else:
                self.observation = rcv_data
                with self._cond:
                    self._cond.notify_all()
